package limit

import (
	"time"
)

type tokenBucket struct {
	size int64
	num  int64
	pre  time.Time
	rate time.Duration //毫秒
}

/**
 * @description: 限流，不支持并发
 * @param {int64} size
 * @param {time.Duration} rate
 * @return {*}
 */
func NewTokenBucket(size int64, rate time.Duration) *tokenBucket {
	return &tokenBucket{
		size: size,
		rate: rate,
		num:  size,
		pre:  time.Now(),
	}
}

/**
 * @description: 获取令牌，会sleep
 * @param {*}
 * @return {*}
 */
func (b *tokenBucket) GetToken() {
	if b.num > 0 {
		b.num--
		return
	}
	t := time.Now()
	if t.Sub(b.pre) < b.rate {
		time.Sleep(b.rate - t.Sub(b.pre))
	}
	b.num = b.size - 1
	b.pre = time.Now()
}
